package service

import (
	"github.com/kordar/k-means-server/src/pojo"
	"github.com/kordar/k-means-server/src/util"
	"log"
)

type BiKMeans struct {
	Data       []pojo.FeatureItem
	FeatureMap map[string][]float64
	Len        int
}

func (k BiKMeans) sse(data []pojo.FeatureItem, center pojo.FeatureItem) float64 {
	var sum float64 = 0
	for _, item := range data {
		if item.Id == center.Id {
			continue
		}
		cosine := util.Cosine(k.FeatureMap[item.Url], k.FeatureMap[center.Url])
		sum = sum + (1 - cosine)
	}
	return sum
}

func getKeys(m map[pojo.FeatureItem][]pojo.FeatureItem) []pojo.FeatureItem {
	j := 0
	keys := make([]pojo.FeatureItem, len(m))
	for k := range m {
		keys[j] = k
		j++
	}
	return keys
}

func (k BiKMeans) compare(result map[pojo.FeatureItem][]pojo.FeatureItem) (min pojo.FeatureItem, max pojo.FeatureItem) {
	keys := getKeys(result)
	sum0 := k.sse(result[keys[0]], keys[0])
	sum1 := k.sse(result[keys[1]], keys[1])
	log.Println("sum ", sum0, sum1, keys[0].Id)
	if sum0 < sum1 {
		return keys[1], keys[0]
	} else {
		return keys[0], keys[1]
	}
}

func (k BiKMeans) getFeatureItemById(id int) *pojo.FeatureItem {
	for _, datum := range k.Data {
		if datum.Id == id {
			return &datum
		}
	}
	return nil
}

func (k BiKMeans) Run() map[pojo.FeatureItem][]pojo.FeatureItem {
	centers := make([]pojo.FeatureItem, 0)
	groups := make(map[pojo.FeatureItem][]pojo.FeatureItem)
	data := k.Data
	for true {
		if len(centers) >= k.Len {
			break
		}
		means := KMeans{Data: data, FeatureMap: k.FeatureMap, Len: 2}
		result := means.Run()
		log.Println("result ===== ", result)
		if len(result) < 2 {
			for item, items := range result {
				groups[item] = items
			}
			break
		}
		min, max := k.compare(result)
		centers = append(centers, max)
		groups[max] = result[max]
		if len(centers)+1 == k.Len {
			centers = append(centers, min)
			groups[min] = result[min]
			break
		}

		data = result[min]
	}

	return groups

}
